package com.zx.sa.service;

import com.zx.sa.common.ModelType;
import jakarta.annotation.Resource;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.multipart.MultipartFile;

@Slf4j
@Service
public class EmbeddingService {


    @Resource(name = "custOllamaAiEmbeddingModel")
    private EmbeddingModel custOllamaAiEmbeddingModel;

    @Resource(name = "custZhiPuAiEmbeddingModel")
    private EmbeddingModel custZhiPuAiEmbeddingModel;

    @Resource(name = "custOpenAiEmbeddingModel")
    private EmbeddingModel custOpenAiEmbeddingModel;

    @Resource
    private VectorStore vectorStore;

    @Resource
    private ChatService chatService;

    private EmbeddingModel getEmbeddingModel() {
        return getEmbeddingModel(ModelType.DEFAULT);
    }

    private EmbeddingModel getEmbeddingModel(String modelType) {
        return switch (ModelType.normalize(modelType)) {
            case ModelType.OPENAI -> custOpenAiEmbeddingModel;
            case ModelType.ZHIPUAI -> custZhiPuAiEmbeddingModel;
            case ModelType.OLLAMA -> custOllamaAiEmbeddingModel;
            default -> custOllamaAiEmbeddingModel;
        };
    }

    public void storeDocuments(List<String> contents) {
        vectorStore.add(contents.stream().map(Document::new).collect(Collectors.toList()));
    }

    public void uploadFile(MultipartFile file) throws IOException {
        String content = IOUtils.toString(file.getInputStream(), StandardCharsets.UTF_8);
        this.storeDocuments(Arrays.asList(content));
    }

    public List<Document> searchData(String query, int topK) {
        // 首先查询向量库
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
                .query(query)
                .topK(topK)
                .build());
        return documents;
    }

    public String embed(String msg, Set<String> fileIds, String modelType) {
        log.debug("embedding... {}", msg);

        Set<String> finalFileIds = (fileIds == null) ? new HashSet<>() : fileIds;

        // 首先查询向量库
        List<Document> documents = searchData(msg, 5);

        // TODO: 引入工具库，更好地识别用户的 prompt
        // TODO: 引入工具库，进行多路召回，并对结果重新进行 Ranking
        String promptContent = documents.stream()
                .filter(doc -> {
                    if (CollectionUtils.isEmpty(finalFileIds)) {
                        return true;
                    }
                    Object fileIdObject = doc.getMetadata().get("file_id");
                    String docFileId = fileIdObject != null ? fileIdObject.toString() : null;
                    return finalFileIds.contains(docFileId);
                })
                .map(Document::getText)
                .filter(StringUtils::hasText)
                .collect(Collectors.joining(" "));

        // 确保 promptContent 不为空
        if (!StringUtils.hasText(promptContent)) {
            promptContent = "No relevant information found in the database.";
        }

        log.debug("Prompt content: {}", promptContent);

        return this.chatService.getChatClient(modelType)
                .prompt(promptContent)  //把 prompt 传递给大模型
                .user(msg)              //用户发送的内容
                .call()
                .content();
    }

    /**
     * 将文本转为向量
     * @param str
     * @return
     */
    public float[] embedding(String str, String modelType) {
        float[] embed = getEmbeddingModel(modelType).embed(str);
        return embed;
    }


}